"""
GroupBy技术
"""
import numpy as np
from pandas import Series, DataFrame
import pandas as pd


def test_01():
    """
    GroupBy
    :return:
    """
    df = DataFrame({'key1': ['a', 'a', 'b', 'b', 'a'],
                    'key2': ['one', 'two', 'one', 'two', 'one'],
                    'data1': np.random.randn(5),
                    'data2': np.random.randn(5)})
    # print(df)
    grouped = df['data1'].groupby(df['key1'])
    # print(grouped)
    # 计算分组平均值 grouped.mean()
    # print(grouped.mean())

    means = df['data1'].groupby([df['key1'], df['key2']]).mean()
    # print(means)
    # print(means.unstack())

    states = np.array(['Ohio', 'California', 'California', 'Ohio', 'Ohio'])
    years = np.array([2005, 2005, 2006, 2005, 2006])
    # print(df['data1'].groupby([states, years]).mean())

    # print(df.groupby('key1').mean())

    # print(df.groupby(['key1', 'key2']).mean())

    print(df.groupby(['key1', 'key2']).size())


def test_02():
    """
    对分组进行迭代
    :return:
    """
    df = DataFrame({'key1': ['a', 'a', 'b', 'b', 'a'],
                    'key2': ['one', 'two', 'one', 'two', 'one'],
                    'data1': np.random.randn(5),
                    'data2': np.random.randn(5)})

    # for name, group in df.groupby('key1'):
    #     print(name)
    #     print(group)

    # for (k1, k2), group in df.groupby(['key1', 'key2']):
    #     print(k1, k2)
    #     print(group)

    # pieces = dict(list(df.groupby('key1')))
    # print(pieces['b'])

    # print(df.dtypes)

    grouped = df.groupby(df.dtypes, axis=1)
    print(dict(list(grouped)))


def test_03():
    """
    选取一个或一组列
    :return:
    """
    df = DataFrame({'key1': ['a', 'a', 'b', 'b', 'a'],
                    'key2': ['one', 'two', 'one', 'two', 'one'],
                    'data1': np.random.randn(5),
                    'data2': np.random.randn(5)})

    # print(df)
    # print(df.groupby('key1')['data1'])
    # print(df.groupby('key1')[['data2']])

    # print(df['data1'].groupby(df['key1']))
    # print(df[['data2']].groupby(df['key1']))

    # print(df.groupby(['key1', 'key2'])[['data2']].mean())
    s_grouped = df.groupby(['key1', 'key2'])['data2']
    print(s_grouped.mean())


def test_04():
    """
    通过字典或Series进行分组
    :return:
    """
    people = DataFrame(np.random.randn(5, 5),
                       columns=['a', 'b', 'c', 'd', 'e'],
                       index=['Joe', 'Steve', 'Wes', 'Jim', 'Travis'])

    # 添加几个NA值
    people.ix[2:3, ['b', 'c']] = np.nan
    print(people)

    mapping = {'a': 'red', 'b': 'red', 'c': 'blue',
               'd': 'blue', 'e': 'red', 'f': 'orange'}

    by_column = people.groupby(mapping, axis=1)
    # print(by_column.sum())

    map_series = Series(mapping)
    print(map_series)

    print(people.groupby(map_series, axis=1).count())


def test_05():
    """
    到270页  通过函数进行分组
    :return:
    """
    pass


def main():
    # test_01()
    # test_02()
    # test_03()
    test_04()


if __name__ == '__main__':
    main()
